# 生成 mcts_games 数据, 需在 dlgo 模块外部的 run.py 中进行调用运行

import argparse
import numpy as np
from .encoder import OnePlaneEncoder
from .board import Board, GameState
from .utils import PrintBoard, PrintMove
from .mcts import MCTSAgent


def GenerateGame(board_size, rounds, max_moves, temperature):
    boards, moves = [], []
    encoder = OnePlaneEncoder(board_size=board_size)
    game = GameState.NewGame(board_size=board_size)
    bot = MCTSAgent(rounds, temperature)
    num_moves = 0
    while not game.IsOver():
        PrintBoard(game.board)
        move = bot.SelectMove(game)
        if move.is_play:
            boards.append(encoder.Encode(game))
            move_one_hot = np.zeros(encoder.NumberPoints())
            move_one_hot[encoder.EncodePoint(move.point)] = 1
            moves.append(move_one_hot)
        PrintMove(game.next_player, move)
        game = game.ApplyMove(move)
        num_moves += 1
        if num_moves > max_moves:
            break
    return np.array(boards), np.array(moves)


def run():
    parser = argparse.ArgumentParser()
    parser.add_argument("--board-size", "-b", type=int, default=9)
    parser.add_argument("--rounds", "-r", type=int, default=1000)
    parser.add_argument("--temperature", "-t", type=float, default=0.8)
    parser.add_argument(
        "--max-moves", "-m", type=int, default=60, help="Max moves per game."
    )
    parser.add_argument("--num-games", "-n", type=int, default=10)
    parser.add_argument("--board-out", default='board_out')
    parser.add_argument("--move-out", default='move_out')

    args = parser.parse_args()  # <1>
    xs = []
    ys = []

    for i in range(args.num_games):
        print('Generating game %d/%d...' % (i + 1, args.num_games))
        x, y = GenerateGame(args.board_size, args.rounds, args.max_moves, args.temperature)  # <2>
        xs.append(x)
        ys.append(y)

    x = np.concatenate(xs)  # <3>
    y = np.concatenate(ys)

    np.save(args.board_out, x)  # <4>
    np.save(args.move_out, y)
